from s4dd import S4forDenovoDesign

# 主程序入口
if __name__ == "__main__":
    # 加载已微调的模型
    finetuned_model_path = "./new_models/s4_length_constrained_iter_5_temp_1.4"
    print(f"加载模型：{finetuned_model_path}")

    # 使用 from_file 方法加载模型
    s4 = S4forDenovoDesign.from_file(finetuned_model_path)
    print("模型加载成功！")

    # 生成分子
    print("开始生成分子...")
    n_designs = 100  # 生成分子的数量
    batch_size = 32  # 生成的批次大小
    temperature = 1.0  # 用于控制生成分子的随机性

    designs, lls = s4.design_molecules(n_designs=n_designs, batch_size=batch_size, temperature=temperature)
    print(f"成功生成 {len(designs)} 个分子！")

    # 将分子以 SMILES 字符串形式保存到文件
    output_file_path = "./molecules.txt"
    with open(output_file_path, "w") as f:
        for i, smiles in enumerate(designs):
            f.write(f"{smiles}\n")
    print(f"生成的分子已保存到 {output_file_path}！")

    # 打印生成的分子（可选）
    print("生成的分子设计如下：")
    for i, smiles in enumerate(designs):
        print(f"{smiles}")
